
from flask import request
from managers.nodes.baseGraphNode import  *
from .sessionInfo import SessionInfo
import json
import uuid
import eventlet
from eventlet import Timeout
from eventlet.hubs import get_hub
from ..execute.baseExeHis import BaseExeHis
import threading


class ToolExecuteNode(BaseGraphNode):
    def __init__(self, id, data, graphNodeManager, random_key):
        super().__init__(id, data, graphNodeManager, random_key)
        self.route = data['data']['route']
        self.timeout = data['data'].get('timeout',10)
        self.flow_next_nodes_failed=[]
        self.flow_next_nodes_success = []
        self.session_pending_events = {}  # {session_id: {correlation_id: event}}
        self.session_pending_hub = {}  # {session_id: {correlation_id: hub_id}}

    def add_flow_next(self, node_id, **kwargs):
        if 'successed' in kwargs['sourceHandle']:
            self.flow_next_nodes_success.append(node_id)
        elif 'failed' in kwargs['sourceHandle']:
            self.flow_next_nodes_failed.append(node_id)
        else:
            self.flow_next_nodes.append(node_id)

    def __del__(self):
        # 清理所有会话的事件
        self.session_pending_events.clear()
        self.session_pending_hub.clear()

    def execute(self, session_info = None, exe_his: BaseExeHis=None, trigger=None, running_id=None):
        if session_info is None: 
            print(f"ToolExecuteNode execute, session_info is None")
            print(f"ToolExecuteNode route: {self.route}")
            socketio_instance = self.graphNodeManager.http_node.get_socketio_instance()

            return


        functionCall = None
        if 'function-call' in self.dict['in']:
            functionFrom = self.dict['in']['function-call']
            prenode =self.graphNodeManager.nodesByIdDict[functionFrom['nodeId']]
            functionCall = prenode.get_newest_output_in_session(session_info, functionFrom['handle'])
        toolName = None
        if 'tool' in self.dict['in']:
            toolFrom = self.dict['in']['tool']
            prenode =self.graphNodeManager.nodesByIdDict[toolFrom['nodeId']]
            if (prenode.type == 'tool'):
                toolName =prenode.get_tool_name()
            else: toolName=None
            # tool = prenode.get_newest_output_in_session(session_info, toolFrom['handle'])

        actureCall=None
        try:
            calls = []
            if(type(functionCall)==str):
                calls = json.loads(functionCall)
            else: calls = functionCall
            for call in calls:
                if call['name'] == toolName:
                    actureCall = call
                    break
        except Exception as e:
            print(e)
            actureCall = None
        with session_info.out_write_lock:
            if session_info.out.get(self.id,None) is None: session_info.out[self.id] = {}
            if('result-msg' not in session_info.out[self.id]): session_info.out[self.id]['result-msg'] = []
            if(functionCall is None):
                session_info.out[self.id]['result-msg'].append("function call not found, check your project edge links")
            if (toolName is None):
                session_info.out[self.id]['result-msg'].append("tool not found, check your project edge links")
            if (actureCall is None):
                session_info.out[self.id]['result-msg'].append("tool name not match, this is normal when agent not pose a matched tool call")

        if (functionCall is None or toolName is None or actureCall is None):
            exe_his.update_content(None)
            if exe_his:exe_his.done()
            self.flow_next(session_info,trigger=trigger, running_id=running_id, no_execute=True)
            return
        
        try:
            # pass
            # self.create_tool_execute_lisener(actureCall, session_info, self.timeout, exe_his,trigger=trigger, running_id=running_id)
            session_info.ws.start_background_task(self.create_tool_execute_lisener,actureCall, session_info, self.timeout, exe_his,trigger=trigger, running_id=running_id)
        except Exception as e:
            print(e)

    def create_tool_execute_lisener(self, actureCall: dict, session_info: SessionInfo, timeout_seconds: int, exe_his, trigger, running_id):
        # 确保该会话的事件字典存在
        if session_info.id not in self.session_pending_events:
            self.session_pending_events[session_info.id] = {}
        if session_info.id not in self.session_pending_hub:
            self.session_pending_hub[session_info.id] = {}
        @session_info.ws.on(self.route, namespace=session_info.full_url)
        def handle_channel_response(msg):
            # print(f"get tool Execute result {msg}")
            received_correlation_id = msg['data']['correlation_id']
            if not received_correlation_id:
                print(f"No correlation_id in message")
                return
            session_events = self.session_pending_events.get(session_info.id, {})
            session_hub = self.session_pending_hub.get(session_info.id, {})
            if received_correlation_id in session_events:
                # print(f"Found matching event for session {session_info.id}, correlation_id {received_correlation_id}")
                pending_event = session_events[received_correlation_id]
                hub = session_hub[received_correlation_id]
                data = msg.get('data')
                # print(f"Sending data to event: {data}")
                # print("thread: event send", threading.current_thread().name)
                # print("hub id: event send", id(eventlet.hubs.get_hub()))
                hub.schedule_call_global(0, pending_event.send, data)
            else:
                print(f"No pending event found for session {session_info.id}, correlation_id {received_correlation_id}")
                    
        correlation_id = str(uuid.uuid4())
        event = eventlet.event.Event()
        # 将事件存储在会话特定的字典中
        self.session_pending_events[session_info.id][correlation_id] = event
        self.session_pending_hub[session_info.id][correlation_id] = get_hub()
        # print(f"Setting up listener for session {session_info.id}, correlation_id: {correlation_id}")
        
        try:
            # print(f"Emitting message for session {session_info.id}, correlation_id: {correlation_id}")
            exe_his.update_content("Waiting for message...")
            session_info.ws.emit(
                self.route, 
                {
                    'data': actureCall, 
                    'correlation_id': correlation_id,
                    'route': self.route
                }, 
                room=session_info.id, 
                namespace=session_info.full_url
            )
            
            with Timeout(timeout_seconds):
                # print("thread: event wait", threading.current_thread().name)
                # print("hub id: event wait", id(eventlet.hubs.get_hub()))
                # print(f"Waiting for response for session {session_info.id} with timeout {timeout_seconds}s")
                response_data = event.wait()
                # print(f"Received response for session {session_info.id}: {response_data}")
                if response_data is None or response_data['data']['success'] == False:
                    success = False
                else:
                    success = True
                response_data = {'session_info': session_info, 'response': response_data['data'], 'success': success}
                exe_his.done()
                self.tool_execution_output(response_data)
                
        except Timeout:
            print(f"Timeout waiting for response for session {session_info.id}")
            self.tool_execution_output({'session_info': session_info, 'response': 'Timeout waiting for response', 'success': False},trigger=trigger, running_id=running_id)
        except Exception as e:
            print(f"Unexpected error for session {session_info.id}: {e}")
            self.tool_execution_output({'session_info': session_info, 'response': 'Unexpected error: ' + str(e), 'success': False},trigger=trigger, running_id=running_id)
        finally:
            # session_info.ws.off(self.route, namespace=session_info.full_url)
            if session_info.id in self.session_pending_events:
                session_events = self.session_pending_events[session_info.id]
                if correlation_id in session_events:
                    del session_events[correlation_id]
                # 如果会话的事件字典为空，也可以清理会话
                if not session_events:
                    del self.session_pending_events[session_info.id]
                
    

    def tool_execution_output(self, response,trigger=None, running_id=None):
        if response is not None:
            success = response.get('success', None)
            session_info = response.get('session_info', None)
            if success:
                with session_info.out_write_lock:
                    if session_info.out.get(self.id,None) is None: session_info.out[self.id] = {}
                    if 'result-msg' not in session_info.out[self.id]: session_info.out[self.id]['result-msg'] = []
                    session_info.out[self.id]['result-msg'].append("Tool "+self.name+" executed successfully with result: "+ str(response['response']))
                self.flow_next(session_info, True, trigger=trigger, running_id=running_id)
            else:
                with session_info.out_write_lock:
                    if session_info.out.get(self.id,None) is None: session_info.out[self.id] = {}
                    if 'result-msg' not in session_info.out[self.id]: session_info.out[self.id]['result-msg'] = []
                session_info.out[self.id]['result-msg'].append("Tool "+self.name+" execution failed")
                self.flow_next(session_info, False, trigger=trigger, running_id=running_id)
        else:
            print("Tool "+self.name+" execution failed")
            with session_info.out_write_lock:
                if session_info.out.get(self.id,None) is None: session_info.out[self.id] = {}
                if 'result-msg' not in session_info.out[self.id]: session_info.out[self.id]['result-msg'] = []
                session_info.out[self.id]['result-msg'].append("Tool "+self.name+" execution failed")
            self.flow_next(session_info, False, trigger=trigger, running_id=running_id)


    def get_all_flow_next_nodes(self):
        return self.flow_next_nodes + self.flow_next_nodes_success + self.flow_next_nodes_failed

    def flow_next(self, session_info: SessionInfo, success:bool = True, no_execute:bool = False, **kwargs):
        if no_execute:
            self.flow_next_process(session_info, self.flow_next_nodes, **kwargs)
        elif success:
            # print("Tool "+self.name+" executed successfully")
            # print(f"next nodes is {self.flow_next_nodes_success}")
            self.flow_next_process(session_info, self.flow_next_nodes_success, **kwargs)
        else:
            self.flow_next_process(session_info, self.flow_next_nodes_failed, **kwargs)
